# 用于测试第一版RFQ自动加载程序，该程序包含自动加载模块、真空判断模块、初始化模块（不包含故障判断）

import time
import configparser
from epics import PV
import re
from datetime import datetime

conf_path = "conf/RFQ.conf"

class RFQ_test(object):
    def __init__(self, pre_para,pv_dict, **kwargs):
        self.init_status = False
        self.powerTarget            = PV(pv_dict['powerTarget'])    # Target Power
        # 初始化参数
        self.openloop_init                  =int(pre_para['openloop_init'])
        self.cavity_amplitude_setpoint_step =int(pre_para['cavity_amplitude_setpoint_step'])
        self.cavity_power_margin            =float(pre_para['cavity_power_margin'])
        self.cavity_vacuum_threshold        =float(pre_para['cavity_vacuum_threshold'])
        # 初始化PV
        self.rf_on                  = PV(pv_dict['rf_on'])
        self.amplitude_cw_drive     = PV(pv_dict["amplitude_cw_drive"])
        self.cavity_power           = PV(pv_dict["cavity_power"])
        self.tracking_frequency_on  = PV(pv_dict["tracking_frequency_on"])
        self.sweep_frequency        = PV(pv_dict["sweep_frequency"])

        self.cavity_power           = PV(pv_dict["cavity_power"],callback= self.on_cavity_power_change)
        self.start_load_on          = PV(pv_dict["start_load_on"])
        self.end_load_on            = PV(pv_dict["end_load_on"])
        self.coupler1_vacuum       =PV(pv_dict["coupler1_vacuum"],callback = self.on_vacuum_change)
        self.coupler2_vacuum       =PV(pv_dict["coupler2_vacuum"],callback = self.on_vacuum_change)
        self.coupler3_vacuum       =PV(pv_dict["coupler3_vacuum"],  callback = self.on_vacuum_change)
        self.coupler4_vacuum       =PV(pv_dict["coupler4_vacuum"], callback= self.on_vacuum_change)
        self.cavity_vacuum         =PV(pv_dict["cavity_vacuum"]  ,callback = self.on_vacuum_change)
        self.Arc_lock              =PV(pv_dict["Arc_lock"]) 
        self.Interlock_status      =PV(pv_dict["Interlock_status"])
        self.RFPWFailStatus1       =PV(pv_dict["RFPWFailStatus1"])
        self.RFPWFailStatus2       =PV(pv_dict["RFPWFailStatus2"])
        self.ResetInterlock        =PV(pv_dict["ResetInterlock"])
        self.ResetPWFaultStat      =PV(pv_dict["ResetPWFaultStat"])
        self.vacuum_status = 0
        self.power_tuning  = None
        self.vacuum_warning_issued = False
        self.interlock   = 0
        
        # 等待所有PV连接
        self.wait_for_pv_connection()
        self.init_status = True

    def wait_for_pv_connection(self):
        """等待所有PV连接成功"""
        pvs = [
            self.powerTarget,
            self.rf_on,
            self.amplitude_cw_drive,
            self.cavity_power,
            self.tracking_frequency_on,
            self.sweep_frequency,
            self.start_load_on,
            self.end_load_on,
            self.coupler1_vacuum,
            self.coupler2_vacuum,
            self.coupler3_vacuum,
            self.coupler4_vacuum,
            self.cavity_vacuum,
            self.Arc_lock,
            self.Interlock_status,
            self.RFPWFailStatus1,
            self.RFPWFailStatus2,
            self.ResetInterlock,
            self.ResetPWFaultStat
        ]
        
        print("等待PV连接...")
        while not all(pv.connected for pv in pvs):
            disconnected_pvs = [pv.pvname for pv in pvs if not pv.connected]
            print(f"等待以下PV连接: {', '.join(disconnected_pvs)}")
            time.sleep(1)
        print("所有PV已连接")

    def check_pv_connection(self):
        """检查所有PV的连接状态"""
        pvs = [
            self.powerTarget,
            self.rf_on,
            self.amplitude_cw_drive,
            self.cavity_power,
            self.tracking_frequency_on,
            self.sweep_frequency,
            self.start_load_on,
            self.end_load_on,
            self.coupler1_vacuum,
            self.coupler2_vacuum,
            self.coupler3_vacuum,
            self.coupler4_vacuum,
            self.cavity_vacuum,
            self.Arc_lock,
            self.Interlock_status,
            self.RFPWFailStatus1,
            self.RFPWFailStatus2,
            self.ResetInterlock,
            self.ResetPWFaultStat
        ]
        
        disconnected_pvs = [pv.pvname for pv in pvs if not pv.connected]
        if disconnected_pvs:
            print(f"警告：以下PV已断开连接: {', '.join(disconnected_pvs)}")
            return False
        return True

    def on_vacuum_change(self, *args, **kwargs):
        try:
            # 检查所有PV是否都已连接
            if not all([
                self.cavity_vacuum.connected,
                self.coupler1_vacuum.connected,
                self.coupler2_vacuum.connected,
                self.coupler3_vacuum.connected,
                self.coupler4_vacuum.connected
            ]):
                print("部分真空PV尚未连接，等待连接...")
                return

            # 检查所有PV是否有有效值
            if any([
                self.cavity_vacuum.value is None,
                self.coupler1_vacuum.value is None,
                self.coupler2_vacuum.value is None,
                self.coupler3_vacuum.value is None,
                self.coupler4_vacuum.value is None
            ]):
                print("部分真空PV值无效，等待有效值...")
                return

            if (self.cavity_vacuum.value < self.cavity_vacuum_threshold and
                self.coupler1_vacuum.value < self.cavity_vacuum_threshold and
                self.coupler2_vacuum.value < self.cavity_vacuum_threshold and
                self.coupler3_vacuum.value < self.cavity_vacuum_threshold and
                self.coupler4_vacuum.value < self.cavity_vacuum_threshold):
                self.vacuum_status = 1
                self.vacuum_warning_issued = False  # 重置警告标志
            else:
                if not self.vacuum_warning_issued:
                    print('vacuum is not ready')
                    self.vacuum_warning_issued = True  # 设置警告标志
                self.vacuum_status = 0
        except Exception as e:
            print(f"真空监控回调发生错误: {str(e)}")
            self.vacuum_status = 0

        
    def on_cavity_power_change(self, *args, **kwargs):
        if self.init_status == True:
            if self.amplitude_cw_drive.value > 10:
                if abs(self.cavity_power.value -self.powerTarget.value) > self.cavity_power_margin:
                    if self.cavity_power.value > self.powerTarget.value:
                        self.power_tuning  = "DECREASE"
                        
                    elif self.cavity_power.value < self.powerTarget.value:
                        self.power_tuning  = "INCREASE"

def read_config(conf_path):
    config = configparser.ConfigParser()
    config.read(conf_path, encoding='utf-8') # 读取配置文件，指定编码为 utf-8
    comment_pattern = re.compile(r'#.*$') # 使用正则表达式匹配所有注释
    
    # 遍历所有的节和键值对，去除值中的注释部分
    for section in config.sections():
        for key, value in config.items(section):
            clean_value = comment_pattern.sub('', value).strip()  # 使用正则表达式移除注释
            config.set(section, key, clean_value)  # 重新设置键值对，去除注释后的值
    print("init llrf config successfully")
    return config

def init_llrf(rfq):
    if rfq.Arc_lock.value == 0 or rfq.Interlock_status.value == 0:
        rfq.interlock = 0
        rfq.ResetInterlock.value = 1
        time.sleep(0.5)
        rfq.ResetInterlock.value = 0
    if rfq.RFPWFailStatus1.value == 0 and rfq.RFPWFailStatus2.value == 0:
        rfq.interlock = 0
        rfq.ResetPWFaultStat.value = 1
        time.sleep(0.5)
        rfq.ResetInterlock.value = 0
    if rfq.Arc_lock.value == 1 and rfq.Interlock_status.value == 1 and rfq.RFPWFailStatus1.value == 1 and rfq.RFPWFailStatus2.value == 1:
        rfq.interlock = 1
    # print('rf_on status:', rfq.rf_on.value,rfq.interlock)
    if rfq.rf_on.value == 0 and rfq.interlock == 1:
       rfq.rf_on.value = 1
       print("start to load rfq power")
    if rfq.vacuum_status == 1 and rfq.rf_on.value == 1:
        if rfq.sweep_frequency == 0:
            rfq.sweep_frequency.value = 1
        if rfq.tracking_frequency_on.value == 0:
            rfq.tracking_frequency_on.value = 1   
        if rfq.amplitude_cw_drive.value < 1 :
            rfq.amplitude_cw_drive.value = rfq.openloop_init
            print("rfq init successfully")

        
    else:
        time.sleep(1)

def auto_load(rfq):
    # 首先检查PV连接状态
    if not rfq.check_pv_connection():
        print("等待PV重新连接...")
        time.sleep(1)
        return
        
    init_llrf(rfq)
    current_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    
    if rfq.power_tuning == "INCREASE" and rfq.rf_on.value == 1 and rfq.vacuum_status == 1:
        rfq.amplitude_cw_drive.value += 10
        print(f"{current_time}: power tuning: INCREASE")
        time.sleep(1)
    elif rfq.power_tuning == "DECREASE" and rfq.rf_on.value == 1 and rfq.vacuum_status == 1:
        # 确保cw_drive值不小于0
        new_value = max(0, rfq.amplitude_cw_drive.value - 5)
        rfq.amplitude_cw_drive.value = new_value
        print(f"{current_time}: power tuning: DECREASE, cw_drive: {new_value}")
        time.sleep(1)
    else:
        time.sleep(1)

def main():
    print("init")
    config = read_config(conf_path)
    pv_dict = config["PVs"]
    pre_para = config["Preset_Parameters"]
    rfq = RFQ_test(pre_para,pv_dict)
    # rfq.on_vacuum_change()
    print("start auto load")
    while True:
        auto_load(rfq)
        time.sleep(1)

if __name__ == '__main__':
    main()


